add weight_sum_fp32 config #2519
Open
+12
−5
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
suexu1025
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
michelle-yooh,
richjames0,
shralex and
vipannalla
as code owners
RissyRan
reviewed
Oct 17, 2025
khatwanimohit
approved these changes
Oct 31, 2025
RissyRan
reviewed
Oct 31, 2025
RissyRan
approved these changes
Oct 31, 2025
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This pull request introduces a new configuration flag, float32_weight_sum, to control the precision of the weighted sum operation in the Mixture of Experts (MoE) layers. The changes are well-implemented and provide useful flexibility for balancing performance and numerical precision.
🔍 General Feedback
- The addition of the
float32_weight_sumflag is a good feature for optimizing MoE layers. - The implementation in
src/MaxText/layers/moe.pycorrectly applies the conditional casting based on the new configuration. - A minor style suggestion was made to improve comment consistency.
| cast_logits_to_fp32: True # whether to cast the logits to fp32. The higher precision is generally beneficial, but it can vary slightly. | ||
| float32_qk_product: False # in dot_product attention, whether to cast to fp32 the inputs to qk product | ||
| float32_logits: False # in dot_product attention, whether to cast to fp32 the inputs to softmax | ||
| float32_weight_sum: True # whether to use full fp32 precision for weight_sum during final unpermute in moe |
There was a problem hiding this comment.
🟢 Nit: For consistency and clarity, it's better to use "MoE" instead of "moe".
Suggested change
| float32_weight_sum: True # whether to use full fp32 precision for weight_sum during final unpermute in moe | |
| float32_weight_sum: True # whether to use full fp32 precision for weight_sum during final unpermute in MoE |
update update Update base.yml Update moe.py Update moe.py
suexu1025
force-pushed
the
qinwen/add_up_quantize_config
branch
from
6f0349d to
ee4e8cc
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Add config flag
weight_sum_fp32for whether to use full fp32 precision for weight_sum during final unpermute in moeTests
final eval loss at 300 steps
2.394 (cloudlog)(https://cloudlogging.app.goo.gl/Q5o2tac9aypGGMyV6)
2.393 (cloudlog)(https://cloudlogging.app.goo.gl/L2N43dAZiHap1Djk7)
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.